import numpy as np


# 定义无人机参数
class UAVParameters:
    def __init__(self):
        self.mass = 1.0  # 质量 (kg)
        self.inertia = np.diag([0.01, 0.01, 0.02])  # 惯性矩阵 (kg·m^2)
        self.gravity = 9.81  # 重力加速度 (m/s^2)
        self.thrust_coefficient = 1e-6
        self.drag_coefficient = 1e-7


# 状态表示 (位置, 速度, 姿态, 角速度)
class State:
    def __init__(self):
        self.position = np.zeros(3)  # x, y, z
        self.velocity = np.zeros(3)  # x_dot, y_dot, z_dot
        self.attitude = np.zeros(3)  # roll, pitch, yaw
        self.angular_velocity = np.zeros(3)  # p, q, r


# 无人机动力学与运动学
class UAVDynamics:
    def __init__(self, params):
        self.params = params
        self.state = State()

    def dynamics(self, state, control_input):
        thrust = control_input[0] * self.params.thrust_coefficient
        torque = control_input[1:4] * self.params.drag_coefficient

        acceleration = np.array([0, 0, -self.params.gravity]) + thrust / self.params.mass
        angular_acceleration = np.linalg.inv(self.params.inertia).dot(torque)

        return acceleration, angular_acceleration

    def update(self, control_input, dt):
        state = self.state

        def f(y):
            pos, vel, att, ang_vel = y[:3], y[3:6], y[6:9], y[9:12]
            state = State()
            state.position, state.velocity, state.attitude, state.angular_velocity = pos, vel, att, ang_vel
            acc, ang_acc = self.dynamics(state, control_input)
            return np.concatenate((vel, acc, ang_vel, ang_acc))

        y = np.concatenate((state.position, state.velocity, state.attitude, state.angular_velocity))

        k1 = f(y)
        k2 = f(y + 0.5 * dt * k1)
        k3 = f(y + 0.5 * dt * k2)
        k4 = f(y + dt * k3)

        y += (dt / 6) * (k1 + 2 * k2 + 2 * k3 + k4)

        self.state.position, self.state.velocity = y[:3], y[3:6]
        self.state.attitude, self.state.angular_velocity = y[6:9], y[9:12]


# 控制器
class UAVController:
    def __init__(self):
        self.position_kp = np.array([1.0, 1.0, 1.0])
        self.position_kd = np.array([0.5, 0.5, 0.5])
        self.attitude_kp = np.array([1.0, 1.0, 1.0])
        self.attitude_kd = np.array([0.5, 0.5, 0.5])

    def compute_control_input(self, state, target_position, target_attitude):
        position_error = target_position - state.position
        velocity_error = -state.velocity
        desired_thrust = self.position_kp * position_error + self.position_kd * velocity_error

        attitude_error = target_attitude - state.attitude
        angular_velocity_error = -state.angular_velocity
        desired_torque = self.attitude_kp * attitude_error + self.attitude_kd * angular_velocity_error

        return np.hstack((desired_thrust, desired_torque))


# 主仿真函数
def run_simulation():
    uav_params = UAVParameters()
    uav_dynamics = UAVDynamics(uav_params)
    controller = UAVController()
    dt = 0.01
    total_time = 5.0
    time = 0.0

    while time < total_time:
        target_position = np.array([0.0, 0.0, 1.0])
        target_attitude = np.array([0.0, 0.0, 0.0])

        control_input = controller.compute_control_input(uav_dynamics.state, target_position, target_attitude)

        uav_dynamics.update(control_input, dt)

        time += dt

        print(f"Time: {time:.2f}s, Position: {uav_dynamics.state.position}, Attitude: {uav_dynamics.state.attitude}")


# 运行仿真
run_simulation()
